from fastapi import APIRouter, Depends, HTTPException, Header
from pydantic import BaseModel
from typing import Optional
import os
import secrets
import asyncio
from datetime import datetime, timedelta
from contextlib import asynccontextmanager

router = APIRouter()

class TokenRequest(BaseModel):
    uid: Optional[str] = None
    limit: Optional[int] = None
    timeOfHours: Optional[int] = None  # 添加timeOfHours参数

# 配置的 API Key
VALID_API_KEY = os.getenv("DOCMEE_API_KEY", "ak_s41kObF3rTp5TEAvtT")

# 存储 token 信息
token_store = {}

# 自动从 API Key 生成 token
def generate_token_from_api_key():
    token = f"sk_{secrets.token_hex(16)}"
    expire_time = datetime.now() + timedelta(hours=2)
    
    token_store[token] = {
        "uid": None,
        "limit": None,
        "expire_time": expire_time
    }
    print(f'生成新 token: {token}')
    return token

# 初始化当前token
current_token = generate_token_from_api_key()

async def check_and_renew_token():
    while True:
        token_info = token_store.get(current_token)
        if not token_info or datetime.now() > token_info["expire_time"]:
            global current_token
            current_token = generate_token_from_api_key()
        await asyncio.sleep(60)  # 每分钟检查一次

@asynccontextmanager
async def lifespan(app: FastAPI):
    # 启动时
    asyncio.create_task(check_and_renew_token())
    yield
    # 关闭时

@router.post("/createApiToken")
async def create_api_token(
    request: TokenRequest,
    api_key: str = Header(None, alias="api-key")
):
    try:
        # 验证 API Key
        if not api_key or api_key != VALID_API_KEY:
            raise HTTPException(
                status_code=401,
                detail="API Key 无效或未提供"
            )

        # 处理过期时间参数
        hours = request.timeOfHours if request.timeOfHours else 2
        valid_hours = min(max(hours, 0), 48)
        expire_time = datetime.now() + timedelta(hours=valid_hours)

        # 生成新token
        token = f"sk_{secrets.token_hex(16)}"

        # 存储token信息
        token_store[token] = {
            "uid": request.uid,
            "limit": request.limit,
            "expire_time": expire_time
        }

        # 如果提供了uid，使之前的token过期
        if request.uid:
            for old_token, info in list(token_store.items()):
                if info["uid"] == request.uid and old_token != token:
                    del token_store[old_token]
                    print(f"uid {request.uid} 的旧 token 已过期: {old_token}")

        print(f'生成新 token: {token} 用户 uid: {request.uid} 限制次数: {request.limit}')

        return {
            "code": 0,
            "message": "操作成功",
            "data": {
                "token": token,
                "expireTime": valid_hours * 3600
            }
        }
    except Exception as e:
        print(f'创建 token 失败: {e}')
        raise HTTPException(
            status_code=500,
            detail=str(e)
        )

@router.get("/currentToken")
async def get_current_token():
    return {
        "code": 0,
        "message": "操作成功",
        "data": {
            "token": current_token
        }
    }
